from typing import Annotated, Any, List, Union
from fastapi import APIRouter, Body
from sqlmodel import select

from api.utils import BaseResponse, get_json_result
from db import SessionDep
from db.constants import LLMType, StatusEnum
from db.model.llm import LLM, LLMFactory, LLMFactoryBase, TenantLLM, TenantLLMBase

llmApi = APIRouter(prefix='/llm', tags=["llm"])

class FactoryItem(LLMFactoryBase):
    model_types: List[str] = []
    llms: List[LLM] = []
    
class FactoriesResponse(BaseResponse):
    data: Union[List[FactoryItem], str, None]

@llmApi.get('/factories', response_model=FactoriesResponse)
async def factories(session: SessionDep) -> Any:
    facs = session.exec(select(LLMFactory)).all()
    facItems = [FactoryItem.model_validate(f) for f in facs \
        if f.status == StatusEnum.VALID.value]
    llms = session.exec(select(LLM)).all()
    mdl_types = {}
    mdl_llms = {}
    for m in llms:
        if m.status != StatusEnum.VALID.value:
            continue
        if m.fid not in mdl_types:
            mdl_types[m.fid] = set([])
        mdl_types[m.fid].add(m.model_type)
        if m.fid not in mdl_llms:
            mdl_llms[m.fid] = []
        mdl_llms[m.fid].append(m)
    for f in facItems:
        f.model_types = list(mdl_types.get(f.name, [LLMType.CHAT, LLMType.EMBEDDING, LLMType.RERANK,
                                                            LLMType.IMAGE2TEXT, LLMType.SPEECH2TEXT, LLMType.TTS]))
        f.llms = mdl_llms.get(f.name, [])
    return get_json_result(data=facItems)

class TenantLLMResponse(BaseResponse):
    data: Union[List[TenantLLMBase], str, None]

@llmApi.get('/tenant/list', response_model=TenantLLMResponse)
async def tenant_llm_list(session: SessionDep) -> Any:
    tenantLlms = session.exec(select(TenantLLM)).all()
    return get_json_result(data=tenantLlms)

@llmApi.post('/tenant/set', response_model=BaseResponse)
async def set_tenant_llm(data: Annotated[TenantLLMBase, Body()], session: SessionDep) -> Any:
    name = data.name
    tenantLlm = session.get(TenantLLM, name)
    if tenantLlm is not None:
        #更新已有数据
        if tenantLlm.api_key != data.api_key:
            data.used_tokens = 0
        llm_data = data.model_dump(exclude_unset=True)
        tenantLlm.sqlmodel_update(llm_data)
        
        # 另一种上流方法：
        # for field in TenantLLMBase.__fields__:
        #     if field != "name":  # 跳过主键字段
        #         setattr(tenantLlm, field, getattr(data, field))
    else:
        #新建数据
        tenantLlm = TenantLLM.model_validate(data)

    session.add(tenantLlm)
    session.commit()
    return get_json_result()